import pandas as pd
import statsmodels.formula.api as smf
from sqlalchemy import create_engine
import pymysql
from matplotlib import pyplot as plt
import seaborn as sns

db_config = {
    "host":'localhost',
    'user':'root',
    'password':'sjk1234',
    'database':'tushare',
    'port':3306,
    'charset':'utf8mb4'
}  
engine = create_engine(f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}")
conn = engine.connect()
chunk_size = 10000

df =pd.read_sql_query("SELECT  d.*, m.buy_lg_vol,m.sell_lg_vol,m.buy_elg_vol, m.sell_elg_vol , m.net_mf_vol,i.vol as i_vol , i.closes as i_closes FROM date_1 d join moneyflows m on d.ts_code = m.ts_code and d.trade_date = m.trade_date LEFT JOIN index_daily i on d.trade_date = i.trade_date and i.ts_code = '399001.SZ' WHERE d.trade_date BETWEEN '2023-01-01' AND '2023-12-31' and d.ts_code = '002229.SZ'",conn,chunksize=chunk_size)

df1 = pd.concat(df,ignore_index=True)
print(df1.head)
df1['zd_closes'] = round((df1['closes'] - df1['closes'].shift(1)) /df1['closes'].shift(1),2)
df1['zs_closes'] = round((df1['i_closes'] - df1['i_closes'].shift(1)) /df1['i_closes'].shift(1),2)
df1['zs_vol'] = round((df1['i_vol'] - df1['i_vol'].shift(1)) /df1['i_vol'].shift(1),2)
print(df1.head)

df1 = df1.dropna(subset=['zd_closes','zs_closes','zs_vol']) 

ex = ['id','trade_date','ts_code','the_date','opens','high','low','closes','pre_closes','changes','pct_chg']
number = df1.select_dtypes(include=['number']).columns.to_list()
number_list = [i for i in number if i not in ex]
formula = 'zd_closes ~ ' + ' + '.join(number_list)
res = smf.ols(formula, data=df1).fit()
print(res.summary())

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

plt.figure(figsize=(10,6))
sns.scatterplot(x='zd_closes',y='zs_closes',data=df1)

plt.figure(figsize=(10,6))
sns.scatterplot(x='zd_closes',y='zs_vol',data=df1)

plt.show()
features = df1[['vol','amount','buy_lg_vol','sell_lg_vol','buy_elg_vol','sell_elg_vol','net_mf_vol','i_vol']]
correlation_matrix = features.corr()
print(correlation_matrix)

plt.figure(figsize=(10,6))
sns.heatmap(correlation_matrix,annot=True,cmap='coolwarm',fmt='.2f',linewidths=.5)
plt.title('Correlation Matrix')
plt.show()
